package main

import (
	"errors"
	"fmt"
	"net"
	"shell/common"
	"shell/common/packages"
	"shell/global"
	"shell/utils"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
)

/**
  代理端需要维护的数据：
  - 被控端保活信息：map，key为uuid
  - 公网转发端口信息：map，key为端口号
**/

var (
	ControlledEndMap = sync.Map{} // key: uuid , value: *ControlledEndInfo
	PortRedirMap     = sync.Map{} // key: uint16 , value: *PortInfo

	ErrPubTcpRedirNotEnable = errors.New("public network TCP port redir is not enable")
	ErrPortOutOfRange       = errors.New("public network TCP redir port is out of range")
)

func ListControlledEnd() []*packages.ControlledStat {
	result := make([]*packages.ControlledStat, 0)
	ControlledEndMap.Range(func(key, value any) bool {
		cei, ok := value.(*ControlledEndInfo)
		if ok && cei.Stat != nil {
			rl := cei.Lock.RLocker()
			rl.Lock()
			cs := *(cei.Stat)
			result = append(result, &cs)
			rl.Unlock()
		}
		return true
	})
	return result
}

func ListPortRedir() []*packages.PortRedirInfo {
	result := make([]*packages.PortRedirInfo, 0)
	PortRedirMap.Range(func(_, value any) bool {
		pi, ok := value.(*PortInfo)
		if ok && pi.RedirRule != nil {
			pri := new(packages.PortRedirInfo)
			pri.UUID = pi.RedirRule.TargetUUID
			pri.RedirDestIp = pi.RedirRule.RedirDestIP
			pri.RedirDestPort = pi.RedirRule.RedirDestPort
			pri.RedirSrcPort = pi.RedirRule.RedirSrcPort
			pri.Stat = packages.PortRedirStatus(pi.Stat.Load())
			pri.ConnNum = pi.ConnNum.Load()
			result = append(result, pri)
		}
		return true
	})
	return result
}

// ControlledEndInfo 被控端信息结构体，如果是被中继节点，则仅记录Stat字段，其他字段都为空
type ControlledEndInfo struct {
	QuicConn      *common.QuicConn         // quic连接
	Stat          *packages.ControlledStat // 在ControlledEndInfo的生命周期中，该字段不能为null，否则报错
	PortRedirRule sync.Map                 // key: uint16 , value: *PortInfo
	Lock          sync.RWMutex             // 用于保护RelayFor切片的读写
	RelayFor      []uuid.UUID
}

func NewControlledEndInfo(stat *packages.ControlledStat, conn *common.QuicConn) (*ControlledEndInfo, error) {
	if stat == nil {
		return nil, global.ErrNil
	}
	if conn != nil && !conn.IsConn.Load() {
		return nil, errors.New("quic conn is down")
	}
	result := new(ControlledEndInfo)
	result.Stat = stat
	result.RelayFor = make([]uuid.UUID, 0)
	result.QuicConn = conn
	return result, nil
}

// OpenStream 开启一个流
func (cei *ControlledEndInfo) OpenStream() (*common.QuicStream, error) {
	if cei.QuicConn == nil || !cei.QuicConn.IsConn.Load() {
		return nil, errors.New("controlled end can't use quic connection")
	}
	return cei.QuicConn.OpenStream()
}

// SyncControlledEnd 同步被控端信息到全局map中
func SyncControlledEnd(req *packages.KeepaliveReq, conn *common.QuicConn) (*ControlledEndInfo, error) {
	if req == nil || req.Stats == nil || len(req.Stats) == 0 || conn == nil {
		return nil, global.ErrNil
	}
	noRelayNum := 0
	var mainNodeInfo *packages.KeepaliveInfo
	relayNode := make([]*packages.KeepaliveInfo, 0)
	relayNodeMap := make(map[uuid.UUID]*packages.KeepaliveInfo)
	for _, v := range req.Stats {
		if v.UseRelay {
			relayNode = append(relayNode, v)
			relayNodeMap[v.UUID] = v
		} else {
			noRelayNum++
			mainNodeInfo = v
		}
	}
	if noRelayNum == 0 || mainNodeInfo == nil {
		return nil, errors.New("error keepalive request, every node use delay")
	}
	if noRelayNum > 1 {
		return nil, errors.New("error keepalive request, more than one master node")
	}
	remoteAddr := conn.Conn.RemoteAddr()
	pubIP, pubPort, _ := utils.ParseIpPort(remoteAddr.String())
	var cei *ControlledEndInfo = nil
	v, have := ControlledEndMap.Load(mainNodeInfo.UUID)
	now := time.Now()
	if have {
		// 节点已经存在，更新部分信息
		cei = (v).(*ControlledEndInfo)
		mainNodeInfo.ToControlledStat(cei.Stat, pubIP, pubPort)
		cei.Stat.LastKeepTime = now
		var relayUUID = make([]uuid.UUID, 0) // 用于存储新请求中中继节点的uuid
		if cei.Stat.IsDelay && len(relayNode) > 0 {
			for _, v := range relayNode {
				if v.IsRelay || !v.UseRelay || v.RelayUUID != mainNodeInfo.UUID {
					cei.Lock.Unlock()
					return nil, errors.New("delay node info error")
				}
				relayUUID = append(relayUUID, v.UUID)
			}
		}
		// 更新代理节点的状态
		cei.Lock.Lock()
		toDelete, toAdd, same := utils.CompareSlice(cei.RelayFor, relayUUID)
		for _, v := range toDelete {
			ControlledEndMap.Delete(v)
		}
		for _, v := range toAdd {
			stat := new(packages.ControlledStat)
			relayNodeMap[v].ToControlledStat(stat, nil, 0)
			stat.LastKeepTime = now
			c, _ := NewControlledEndInfo(stat, nil)
			ControlledEndMap.Store(v, c)
		}
		for _, v := range same {
			// 是已有的，需要更新
			value, ok := ControlledEndMap.Load(v)
			cNew := relayNodeMap[v]
			if ok {
				cOld := value.(*ControlledEndInfo)
				cOld.Stat.LastKeepTime = now
				cNew.ToControlledStat(cOld.Stat, nil, 0)
			} else {
				stat := new(packages.ControlledStat)
				stat.LastKeepTime = now
				cNew.ToControlledStat(stat, nil, 0)
				c, _ := NewControlledEndInfo(stat, nil)
				ControlledEndMap.Store(v, c)
			}
		}
		cei.RelayFor = relayUUID
		cei.Lock.Unlock()
		return cei, nil
	} else {
		// 节点不存在，新建节点
		stat := new(packages.ControlledStat)
		stat.LastKeepTime = now
		mainNodeInfo.ToControlledStat(stat, pubIP, pubPort)
		cei, _ = NewControlledEndInfo(stat, conn)
		var relayUUID []uuid.UUID
		if cei.Stat.IsDelay && len(relayNode) > 0 {
			for _, v := range relayNode {
				if v.IsRelay || !v.UseRelay || v.RelayUUID != mainNodeInfo.UUID {
					cei.Lock.Unlock()
					return nil, errors.New("delay node info error")
				}
				relayUUID = append(relayUUID, v.UUID)
			}
		}
		// 添加代理节点
		cei.Lock.Lock()
		for _, v := range relayUUID {
			k := relayNodeMap[v]
			cs := new(packages.ControlledStat)
			cs.LastKeepTime = now
			c, _ := NewControlledEndInfo(cs, nil)
			k.ToControlledStat(c.Stat, nil, 0)
			ControlledEndMap.Store(v, c)
		}
		cei.Lock.Unlock()
		ControlledEndMap.Store(mainNodeInfo.UUID, cei)
		go cei.WaitClose()
		return cei, nil
	}
}

func (cei *ControlledEndInfo) IsProxy() bool {
	if cei.Stat == nil {
		return false
	}
	return cei.Stat.UUID == globalCfg.UUID
}

// WaitClose 一旦保活连接断开，则关闭quic连接、在map中删除相关信息，关闭相关端口转发
func (cei *ControlledEndInfo) WaitClose() {
	<-cei.QuicConn.Ctx().Done()
	cei.Lock.Lock()
	// 从map中删除相关信息
	ControlledEndMap.Delete(cei.Stat.UUID)
	// 从map中删除中继节点信息
	if len(cei.RelayFor) > 0 {
		for _, v := range cei.RelayFor {
			ControlledEndMap.Delete(v)
		}
	}
	cei.Lock.Unlock()
	// 关闭端口转发
	cei.PortRedirRule.Range(func(_, value any) bool {
		pi, ok := value.(*PortInfo)
		if ok {
			_ = pi.HardStop()
		}
		return true
	})
}

type PortInfo struct {
	Port             uint16
	RedirRule        *packages.PortRedirReq
	TcpListener      net.Listener
	Stat             atomic.Int32 // 指示端口的状态
	CanAcceptNewConn atomic.Bool  // 指示端口是否能接受新的连接
	ConnNum          atomic.Int32 // 指示连接数
}

// GetOrCreatePortInfo 新建一个端口信息，并添加到端口信息映射中
func GetOrCreatePortInfo(port uint16) (*PortInfo, error) {
	if globalCfg.TcpPorts == nil {
		return nil, ErrPubTcpRedirNotEnable
	}
	if globalCfg.TcpPorts.From > port || globalCfg.TcpPorts.To < port {
		return nil, ErrPortOutOfRange
	}
	value, ok := PortRedirMap.Load(port)
	if !ok {
		result := new(PortInfo)
		result.Port = port
		result.CanAcceptNewConn.Store(false)
		result.Stat.Store(int32(packages.PortHardStop))
		result.RedirRule = nil
		result.TcpListener = nil
		PortRedirMap.Store(port, result)
		return result, nil
	} else {
		pi := value.(*PortInfo)
		return pi, nil
	}
}

// SoftStop 软关闭，仅修改状态
func (pi *PortInfo) SoftStop() error {
	pi.CanAcceptNewConn.Store(false)
	pi.Stat.Store(int32(packages.PortSoftStop))
	return nil
}

// HardStop 硬关闭，需要关闭tcp端口监听
func (pi *PortInfo) HardStop() error {
	pi.CanAcceptNewConn.Store(false)
	pi.Stat.Store(int32(packages.PortHardStop))
	_ = pi.TcpListener.Close()
	pi.TcpListener = nil
	pi.ConnNum.Store(0)
	return nil
}

// Start 启动端口转发，不阻塞
func (pi *PortInfo) Start(rule *packages.PortRedirReq) error {
	// 校验规则是否符合开启条件
	if rule == nil {
		return global.ErrNil
	}
	checkRule, err := rule.CheckAndClone(true, Role)
	if err != nil {
		return err
	}
	pi.RedirRule = checkRule
	// 开启监听，并处理转发请求
	if !pi.Stat.CompareAndSwap(int32(packages.PortHardStop), int32(packages.PortOpen)) {
		return errors.New("current port status is not hard stop, can't start")
	}
	if pi.TcpListener != nil {
		_ = pi.TcpListener.Close()
		pi.TcpListener = nil
	}
	value, have := ControlledEndMap.Load(pi.RedirRule.TargetUUID)
	if !have {
		return errors.New("target node is down")
	}
	cei := value.(*ControlledEndInfo)
	if cei.IsProxy() {
		return errors.New("can't set target node to proxy end")
	}
	cei.PortRedirRule.Store(pi.Port, pi)
	l, e := net.Listen("tcp", net.JoinHostPort("0.0.0.0", fmt.Sprintf("%d", pi.Port)))
	if e != nil {
		pi.Stat.Store(int32(packages.PortHardStop))
		return e
	}
	pi.TcpListener = l
	pi.CanAcceptNewConn.Store(true)
	go func(portInfo *PortInfo, cei *ControlledEndInfo) {
		for {
			conn, err := portInfo.TcpListener.Accept()
			if err != nil {
				logger.Infof(0, "tcp error accept connection: %v", err)
				break
			}
			if portInfo.RedirRule == nil {
				_ = conn.Close()
				break
			}
			if !portInfo.CanAcceptNewConn.Load() {
				_ = conn.Close()
				continue
			}
			go portInfo.HandlePortRedir(conn, cei.QuicConn)
		}
		_ = portInfo.HardStop()
	}(pi, cei)
	return nil
}

// HandlePortRedir 处理一个端口转发连接
func (pi *PortInfo) HandlePortRedir(tcpConn net.Conn, quicConn *common.QuicConn) {
	stream, err := quicConn.OpenStream()
	if err != nil {
		common.HandleErr(err)
		_ = tcpConn.Close()
		return
	}
	err = stream.WritePackage(pi.RedirRule)
	if err != nil {
		common.HandleErr(err)
		_ = tcpConn.Close()
		return
	}
	resp := new(packages.PortRedirResp)
	err = stream.ReadPackage(resp)
	if err != nil {
		common.HandleErr(err)
		_ = tcpConn.Close()
		return
	}
	if err = resp.CompareFromReq(pi.RedirRule); err != nil {
		logger.Errorf(0, "error redir port: %v", err)
		_ = tcpConn.Close()
		_ = stream.Close()
		return
	}
	if resp.GetRole() != packages.RoleControlled {
		logger.Errorf(0, "error resp role")
		_ = tcpConn.Close()
		_ = stream.Close()
		return
	}
	if !resp.IsOk {
		logger.Infof(0, "error redir port: %s", string(resp.ErrorInfo))
		_ = tcpConn.Close()
		_ = stream.Close()
		return
	}
	pi.ConnNum.Add(1)
	common.NewTimeoutRWCPairWarp(stream, tcpConn, global.DefaultIdleTime, timeoutWatcher).RedirSyncNoKey()
	pi.ConnNum.Add(-1)
	if pi.ConnNum.CompareAndSwap(0, 0) && pi.Stat.CompareAndSwap(int32(packages.PortSoftStop), int32(packages.PortHardStop)) {
		// 如果当前状态是软关闭且连接数为0时，转换为硬关闭状态
		_ = pi.HardStop()
	}
}
